// 方法: getUnreadBadgeCount()
import { screen, imageResource, FileType } from '@nut-tree/nut-js';
import '@nut-tree/template-matcher';
import Jimp from 'jimp';
import cv from 'opencv4nodejs-prebuilt';
import Tesseract from 'tesseract.js';
import { mkdir } from 'fs/promises';
import path from 'path';

screen.config.resourceDirectory = './src/wechat/templates';
screen.config.autoHighlight = true;
screen.config.highlightDurationMs = 800;

async function captureScreenshot(screenshotDir) {
  await mkdir(screenshotDir, { recursive: true });
  const filename = `wechat-${Date.now()}.png`;
  await screen.capture(filename, FileType.PNG, screenshotDir);
  return path.resolve(screenshotDir, filename);
}

async function ocrDigitsFromBuffer(buffer) {
  const result = await Tesseract.recognize(buffer, 'eng', {
    tessedit_char_whitelist: '0123456789',
  });
  const text = (result.data.text || '').replace(/\D+/g, '');
  return text ? parseInt(text, 10) : 0;
}

function pickLargestRedContourRect(matBgr) {
  const matHsv = matBgr.cvtColor(cv.COLOR_BGR2HSV);
  // 红色在 HSV 里通常跨两个区间
  const lower1 = new cv.Vec(0, 100, 100);
  const upper1 = new cv.Vec(10, 255, 255);
  const lower2 = new cv.Vec(160, 100, 100);
  const upper2 = new cv.Vec(179, 255, 255);
  const mask1 = matHsv.inRange(lower1, upper1);
  const mask2 = matHsv.inRange(lower2, upper2);
  const mask = mask1.bitwiseOr(mask2);

  const kernel = cv.getStructuringElement(cv.MORPH_ELLIPSE, new cv.Size(3, 3));
  const morphed = mask.morphologyEx(kernel, cv.MORPH_CLOSE);

  const contours = morphed.findContours(cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE);
  if (!contours.length) return null;

  // 选面积最大的红色区域
  const largest = contours.reduce((acc, cur) => (cur.area > acc.area ? cur : acc));
  const rect = largest.boundingRect();
  return rect;
}

export default async function getUnreadBadgeCount() {
  // Step 1: 截屏
  const screenshotPath = await captureScreenshot('./.tmp');

  // Step 2: 通过“消息”模板定位 Badge 大致位置（主方案）
  let roiRect = null;
  try {
    const msgRegion = await screen.find(imageResource('消息-template.png'), { confidence: 0.7 });
    screen.highlight(msgRegion);

    // 假设 Badge 在“消息”文本右侧偏上，给一个可调的 ROI（偏移与尺寸可按需微调）
    const x = Math.max(0, msgRegion.left + msgRegion.width + 8);
    const y = Math.max(0, msgRegion.top - 4);
    const w = 48; // 宽度足够覆盖红色圆和数字
    const h = 28; // 高度足够覆盖红色圆和数字
    roiRect = new cv.Rect(x, y, w, h);
  } catch (e) {
    // 模板匹配失败时不报错，走兜底方案
  }

  // Step 3: 读取截图并确定 ROI（若主方案失败，兜底用颜色阈值找红色区域）
  const mat = cv.imread(screenshotPath); // BGR
  let roiMat = null;

  if (roiRect) {
    // 直接使用模板匹配得到的 ROI
    const safeRect = new cv.Rect(
      roiRect.x,
      roiRect.y,
      Math.min(roiRect.width, mat.cols - roiRect.x),
      Math.min(roiRect.height, mat.rows - roiRect.y)
    );
    roiMat = mat.getRegion(safeRect);
  } else {
    // 颜色兜底：在整图中找最大的红色区域作为 Badge
    const rect = pickLargestRedContourRect(mat);
    if (!rect) {
      console.log('未找到红色 Badge');
      return 0;
    }
    // 适当扩展一点边界，方便 OCR
    const pad = 4;
    const x = Math.max(0, rect.x - pad);
    const y = Math.max(0, rect.y - pad);
    const w = Math.min(mat.cols - x, rect.width + pad * 2);
    const h = Math.min(mat.rows - y, rect.height + pad * 2);
    roiMat = mat.getRegion(new cv.Rect(x, y, w, h));
  }

  // Step 4: 预处理以提高 OCR 识别率（灰度 + 二值 + 轻微膨胀）
  let gray = roiMat.cvtColor(cv.COLOR_BGR2GRAY);
  gray = gray.threshold(0, 255, cv.THRESH_BINARY + cv.THRESH_OTSU);
  const kernel = cv.getStructuringElement(cv.MORPH_RECT, new cv.Size(2, 2));
  gray = gray.dilate(kernel);

  // Step 5: OCR（仅识别数字）
  const pngBuffer = cv.imencode('.png', gray); // 直接传入 Buffer
  const count = await ocrDigitsFromBuffer(pngBuffer);

  console.log('未读消息数:', count);
  return count;
}

// 自运行（保持你原来的入口）
getUnreadBadgeCount();